CompositeKeyAlreadyAssignedIdentifierInsertionManager.java

package org.codefilarete.stalactite.mapping.id.manager;

import java.util.Set;
import java.util.function.Consumer;
import java.util.function.Function;

import org.codefilarete.stalactite.engine.listener.InsertListener;
import org.codefilarete.stalactite.engine.listener.PersistListener;
import org.codefilarete.stalactite.engine.listener.SelectListener;
import org.codefilarete.tool.collection.Collections;

/**
 *
 * @param <C>
 * @param <I>
 * @author Guillaume Mary
 */
public class CompositeKeyAlreadyAssignedIdentifierInsertionManager<C, I> extends AlreadyAssignedIdentifierManager<C, I> {
	
	private final ThreadLocal<Set<C>> currentMarkedPersistedEntities;
	
	private final CurrentContextManagerInsertListener currentContextManagerInsertListener = new CurrentContextManagerInsertListener();
	
	private final CurrentContextManagerSelectListener currentContextManagerSelectListener = new CurrentContextManagerSelectListener();
	
	private final CurrentContextManagerPersistListener currentContextManagerPersistListener = new CurrentContextManagerPersistListener();
	
	public CompositeKeyAlreadyAssignedIdentifierInsertionManager(Class<I> identifierType,
																 Consumer<C> markAsPersistedFunction,
																 Function<C, Boolean> isPersistedFunction) {
		super(identifierType, markAsPersistedFunction, isPersistedFunction);
		this.currentMarkedPersistedEntities = null;
	}
	
	public CompositeKeyAlreadyAssignedIdentifierInsertionManager(Class<I> identifierType) {
		this(identifierType, new ThreadLocal<>());
	}
	
	private CompositeKeyAlreadyAssignedIdentifierInsertionManager(Class<I> identifierType, ThreadLocal<Set<C>> currentMarkedPersistedEntities) {
		super(identifierType, o -> currentMarkedPersistedEntities.get().add(o), o -> currentMarkedPersistedEntities.get().contains(o));
		this.currentMarkedPersistedEntities = currentMarkedPersistedEntities;
	}
	
	public PersistListener<C> getPersistListener() {
		return currentContextManagerPersistListener;
	}
	
	@Override
	public InsertListener<C> getInsertListener() {
		return currentContextManagerInsertListener;
	}
	
	@Override
	public SelectListener<C, I> getSelectListener() {
		return currentContextManagerSelectListener;
	}
	
	
	private void ensureCurrentEntitySet() {
		currentMarkedPersistedEntities.set(Collections.newIdentitySet(100));
	}
	
	private void clearCurrentEntitySet() {
		currentMarkedPersistedEntities.remove();
	}
	
	private class CurrentContextManagerPersistListener implements PersistListener<C> {
		
		@Override
		public void beforePersist(Iterable<? extends C> entities) {
			ensureCurrentEntitySet();
		}
		
		@Override
		public void afterPersist(Iterable<? extends C> entities) {
			clearCurrentEntitySet();
		}
	}
	
	private class CurrentContextManagerInsertListener implements InsertListener<C> {
		
		@Override
		public void beforeInsert(Iterable<? extends C> entities) {
			ensureCurrentEntitySet();
		}
		
		@Override
		public void afterInsert(Iterable<? extends C> entities) {
			markAsPersisted(entities);
			clearCurrentEntitySet();
		}
		
		@Override
		public void onInsertError(Iterable<? extends C> entities, RuntimeException runtimeException) {
			clearCurrentEntitySet();
		}
	}
	
	private class CurrentContextManagerSelectListener implements SelectListener<C, I> {
		
		@Override
		public void beforeSelect(Iterable<I> ids) {
			ensureCurrentEntitySet();
		}
		
		@Override
		public void afterSelect(Set<? extends C> result) {
			markAsPersisted(result);
			clearCurrentEntitySet();
		}
		
		@Override
		public void onSelectError(Iterable<I> ids, RuntimeException exception) {
			clearCurrentEntitySet();
		}
	}
}